library(survival)
library(statmod)
library(nleqslv)
library(rootSolve)
library(optimx)

options(width=200)
rm(list = ls())
graphics.off()

n<-1000
b0<--1
beta0<-c(1,1)
qq<-1/2
samp<-1000
seed0 = seq(1,samp)+118171
bcen<-0.4
#bcen<-0.17
censoring<-0
bw<-0.01
cntl = list(1e-6,1e-6,1e-3,500)
names(cntl) = c('xtol','ftol','btol','maxit')
bsam<-500

eb0<-vector()
ebeta1<-matrix(0,samp,3)
jac<-matrix(0,3,3)
iter<-vector()
seboots<-matrix(0,samp,3)
ci1<-matrix(0,samp,3)
ci2<-matrix(0,samp,3)

# U(beta) - for finding the root
sn2<-function(b,z1,z2,w1,w2,h1,h2,y,del,km){
  int1<-(y<w1)*y
  int2<-(w1<y)*(y<w2)*(w1+(y-w1)*exp(b[2]*h1))
  int3<-(w2<y)*(w1+exp(b[2]*h1)*(w2-w1)+exp(b[3]*h2)*(y-w2))
  integ<-int1+int2+int3
  integ<-exp(b[1])*integ
  a<-1/(1+exp(-(integ-1)/bw)) # logistic function as smoothing function
  temp0<-del/km*(a-qq)
  temp1<-del/km*(a-qq)*z1
  temp2<-del/km*(a-qq)*z2
  fun<-(c(mean(temp0,na.rm=TRUE),mean(temp1,na.rm=TRUE),mean(temp2,na.rm=TRUE)))
  return(fun)
}
jacobian<-function(b,z1,z2,w1,w2,h1,h2,y,del,km){
  int1<-(y<w1)*y
  int2<-(w1<y)*(y<w2)*(w1+(y-w1)*exp(b[2]*h1))
  int3<-(w2<y)*(w1+exp(b[2]*h1)*(w2-w1)+exp(b[3]*h2)*(y-w2))
  integ<-int1+int2+int3
  integ<-exp(b[1])*integ
  a<-1/bw*exp(-(integ-1)/bw)/(1+exp(-(integ-1)/bw))^2
  int21<-(w1<y)*(y<w2)*((y-w1)*exp(b[2]*h1)*h1)
  int31<-(w2<y)*(exp(b[2]*h1)*(w2-w1)*h1)
  int32<-(w2<y)*(exp(b[3]*h2)*(y-w2)*h2)
  jac[1,1]<-mean(del/km*a*integ)
  jac[1,2]<-mean(del/km*a*(int21+int31)*exp(b[1]))
  jac[1,3]<-mean(del/km*a*(int32)*exp(b[1]))
  jac[2,1]<-mean(del/km*a*z1*integ)
  jac[2,2]<-mean(del/km*a*z1*(int21+int31)*exp(b[1]))
  jac[2,3]<-mean(del/km*a*z1*(int32)*exp(b[1]))
  jac[3,1]<-mean(del/km*a*z2*integ)
  jac[3,2]<-mean(del/km*a*z2*(int21+int31)*exp(b[1]))
  jac[3,3]<-mean(del/km*a*z2*(int32)*exp(b[1]))
  return(jac)
}
# weighted version of U(beta)
wsn2<-function(b,weight,z1,z2,w1,w2,h1,h2,y,del,km){
  int1<-(y<w1)*y
  int2<-(w1<y)*(y<w2)*(w1+(y-w1)*exp(b[2]*h1))
  int3<-(w2<y)*(w1+exp(b[2]*h1)*(w2-w1)+exp(b[3]*h2)*(y-w2))
  integ<-int1+int2+int3
  integ<-exp(b[1])*integ
  a<-1/(1+exp(-(integ-1)/bw)) # logistic function as smoothing function
  temp0<-weight*del/km*(a-qq)
  temp1<-weight*del/km*(a-qq)*z1
  temp2<-weight*del/km*(a-qq)*z2
  fun<-(c(mean(temp0,na.rm=TRUE),mean(temp1,na.rm=TRUE),mean(temp2,na.rm=TRUE)))
  return(fun)
}
wjacobian<-function(b,weight,z1,z2,w1,w2,h1,h2,y,del,km){
  int1<-(y<w1)*y
  int2<-(w1<y)*(y<w2)*(w1+(y-w1)*exp(b[2]*h1))
  int3<-(w2<y)*(w1+exp(b[2]*h1)*(w2-w1)+exp(b[3]*h2)*(y-w2))
  integ<-int1+int2+int3
  integ<-exp(b[1])*integ
  a<-1/bw*exp(-(integ-1)/bw)/(1+exp(-(integ-1)/bw))^2
  int21<-(w1<y)*(y<w2)*((y-w1)*exp(b[2]*h1)*h1)
  int31<-(w2<y)*(exp(b[2]*h1)*(w2-w1)*h1)
  int32<-(w2<y)*(exp(b[3]*h2)*(y-w2)*h2)
  jac[1,1]<-mean(weight*del/km*a*integ)
  jac[1,2]<-mean(weight*del/km*a*(int21+int31)*exp(b[1]))
  jac[1,3]<-mean(weight*del/km*a*(int32)*exp(b[1]))
  jac[2,1]<-mean(weight*del/km*a*z1*integ)
  jac[2,2]<-mean(weight*del/km*a*z1*(int21+int31)*exp(b[1]))
  jac[2,3]<-mean(weight*del/km*a*z1*(int32)*exp(b[1]))
  jac[3,1]<-mean(weight*del/km*a*z2*integ)
  jac[3,2]<-mean(weight*del/km*a*z2*(int21+int31)*exp(b[1]))
  jac[3,3]<-mean(weight*del/km*a*z2*(int32)*exp(b[1]))
  return(jac)
}
for (i in 1:samp){
  set.seed(seed0[i])
  #sampling#
  z1<-rexp(n,1) #instrumental vaiable 1
  z2<-rexp(n,1) # instrumental variable 2
  w1<-rexp(n,4)
  w2<-w1+rexp(n,4)   #w1 is always < w2
  #  w1<-0.6
  #  w2<-0.9
  h1<-z1+rgamma(n,4,scale=0.4)
  h2<-z2+rgamma(n,4,scale=0.4)
  h1<-h1/2
  h2<-h2/2
  tau<-rgamma(n,h1,1)/qgamma(0.5,h1,1)
  tau<-exp(-b0)*tau
  
  ind1<-(tau<w1)*tau
  temp2<-w1+(w2-w1)*exp(beta0[1]*h1)
  ind2<-(w1<=tau)*(tau<temp2)*(w1+(tau-w1)/exp(beta0[1]*h1))
  ind3<-(tau>=temp2)*(w2+exp(-beta0[2]*h2)*(tau-temp2))
  t<-ind1+ind2+ind3
  c<-rexp(n,bcen)
  y<-ifelse(t<=c,t,c)
  del<-(t<=c)*1
  cdel<-1-del
  censoring<-censoring+1-mean(del)
  
  #Kaplan-Meier of censoring
  ranky<-rank(y)
  fit<-survfit(Surv(y,cdel)~1)
  if (fit$surv[n]==0) fit$surv[n]=fit$surv[n-1]
  km<-fit$surv[ranky]
  
  ebeta1[i,]<-c(-0.5,1/2,1/2) # initial values
  est = ebeta1[i,]
  jactry<-jacobian(est,z1,z2,w1,w2,h1,h2,y,del,km)
  
  funest<-nleqslv(est,sn2,jac=NULL,method="Broyden",z1=z1,z2=z2,w1=w1,w2=w2,h1=h1,h2=h2,y=y,del=del,km=km,control=cntl)
  ebeta1[i,]<-funest$x[1:3]
  if ((abs(ebeta1[i,1])>7) | (abs(ebeta1[i,2])>7) | (abs(ebeta1[i,3])>7)) {
    ebeta1[i,]<-c(NA,NA,NA);
  }
  print(c(i,ebeta1[i,]))
  
  #######################################
  ##Bootstrap SE and confidence intervals
  bebeta1<-matrix(0,bsam,3)
  for (b in 1:bsam) {
    weight<-rexp(n,1)
    weight<-weight/mean(weight)
    #Kaplan-Meier of censoring
    ranky<-rank(y) 
    fit<-survfit(Surv(y,cdel)~1,weights=weight)
    tempn<-length(fit$surv)
    if (fit$surv[tempn]==0) fit$surv[tempn]=fit$surv[tempn-1]
    km<-fit$surv[ranky]
    
    bebeta1[b,]<-c(-1,1,1) # initial values
    est = bebeta1[b,]
    funest<-nleqslv(est,wsn2,jac=NULL,method="Broyden",weight=weight,z1=z1,z2=z2,w1=w1,w2=w2,h1=h1,h2=h2,y=y,del=del,km=km,control=cntl)
    bebeta1[b,]<-funest$x[1:3]
    if ((abs(bebeta1[b,1])>5) | (abs(bebeta1[b,2])>5) | (abs(bebeta1[b,3])>5)) {
      bebeta1[b,]<-c(NA,NA,NA);
    }
  } #end bootstrap sam
  #######################################
  seboots[i,]<-c(sqrt(var(bebeta1[,1],na.rm = TRUE)),sqrt(var(bebeta1[,2],na.rm = TRUE)),sqrt(var(bebeta1[,3],na.rm = TRUE)))
  temp1<-((ebeta1[i,1]-1.96*seboots[i,1])<=b0)*((ebeta1[i,1]+1.96*seboots[i,1])>=b0)
  temp2<-((ebeta1[i,2]-1.96*seboots[i,2])<=beta0[1])*((ebeta1[i,2]+1.96*seboots[i,2])>=beta0[1])
  temp3<-((ebeta1[i,3]-1.96*seboots[i,3])<=beta0[2])*((ebeta1[i,3]+1.96*seboots[i,3])>=beta0[2])
  ci1[i,]<-c(temp1,temp2,temp3)
  q1<-quantile(bebeta1[,1],prob=c(0.025,0.975),na.rm = TRUE); temp1<-(q1[1]<=b0)*(q1[2]>=b0)
  q1<-quantile(bebeta1[,2],prob=c(0.025,0.975),na.rm = TRUE); temp2<-(q1[1]<=beta0[1])*(q1[2]>=beta0[1])
  q1<-quantile(bebeta1[,3],prob=c(0.025,0.975),na.rm = TRUE); temp3<-(q1[1]<=beta0[2])*(q1[2]>=beta0[2])
  ci2[i,]<-c(temp1,temp2,temp3)
  print(c(i,seboots[i,],ci1[i,],ci2[i,]))
  
} #end main loop of i